from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution
from launch.conditions import IfCondition, UnlessCondition
from launch_ros.substitutions import FindPackageShare
from launch_ros.actions import Node
from ament_index_python.packages import get_package_share_directory

import os

def generate_launch_description():
    """
    Generates the launch description for the ZOROS robot, with a more elegant
    and readable structure for conditional GUI launching.
    """

    # --- 1. Declare Launch Arguments ---
    # These arguments can be set from the command line (e.g., `ros2 launch ... gui:=true`)
    
    # Argument to specify the robot model
    model_arg = DeclareLaunchArgument(
        'model', 
        default_value='zoram_one_alpha',
        description='The model name of the robot.'
    )

    # Argument to switch between simulation and real hardware modes
    simulation_arg = DeclareLaunchArgument(
        'simulation', 
        default_value='true',
        description='Flag to run in simulation mode.'
    )

    # Argument to conditionally launch the joint control GUI
    gui_arg = DeclareLaunchArgument(
        'gui', 
        default_value='false',
        description='Flag to launch the Joint Control GUI.'
    )

    # --- 2. Prepare Paths and Configurations ---

    # Construct the full path to the URDF file
    urdf_path = os.path.join(
        get_package_share_directory('zoram_one_alpha'),
        'urdf',
        'zoram_one_alpha.urdf'
    )
    
    # Read the URDF file content for the robot_state_publisher
    with open(urdf_path, 'r') as urdf_file:
        robot_description_content = urdf_file.read()

    # Construct the full path to the RViz configuration file
    rviz_config_path = PathJoinSubstitution([
        FindPackageShare('zoram_one_alpha'),
        'rviz',
        'dual_robot_preview.rviz'
    ])

    # Ghost robot URDF for preview
    ghost_urdf_path = os.path.join(
        get_package_share_directory('zoram_one_alpha'),
        'urdf',
        'zoram_one_alpha_ghost.urdf'
    )
    with open(ghost_urdf_path, 'r') as urdf_file:
        ghost_robot_description_content = urdf_file.read()

    # --- 3. Define Node Actions ---

    # Node for the main robot controller
    robot_manager_node = Node(
        package='zoros',
        executable='robot_manager_node',
        name='robot_manager_node',
        parameters=[{'is_simulation': LaunchConfiguration('simulation')}],
        output='screen'
    )

    # Node to publish robot model states to TF (real robot)
    robot_state_publisher_node = Node(
        package='robot_state_publisher',
        executable='robot_state_publisher',
        name='robot_state_publisher',
        parameters=[{'robot_description': robot_description_content}],
        output='screen'
    )

    # Ghost preview components only when GUI (and thus preview) is enabled.
    static_tf_ghost = Node(
        package='tf2_ros',
        executable='static_transform_publisher',
        name='ghost_base_static_tf',
        arguments=['0','0','0','0','0','0','base_link','ghost_base_footprint'],
        output='screen',
        condition=IfCondition(LaunchConfiguration('gui'))
    )

    ghost_robot_state_publisher_node = Node(
        package='robot_state_publisher',
        executable='robot_state_publisher',
        name='ghost_robot_state_publisher',
        parameters=[{'robot_description': ghost_robot_description_content}],
        remappings=[
            ('joint_states', '/preview_joint_states'),
            ('/robot_description', '/robot_description_ghost')
        ],
        output='screen',
        condition=IfCondition(LaunchConfiguration('gui'))
    )

    # RViz node WITH ghost (when gui true)
    rviz_node_with_ghost = Node(
        package='rviz2',
        executable='rviz2',
        name='rviz2',
        arguments=['-d', rviz_config_path],
        parameters=[
            {'robot_description': robot_description_content},
            {'robot_description_ghost': ghost_robot_description_content}
        ],
        output='screen',
        condition=IfCondition(LaunchConfiguration('gui'))
    )

    # RViz node WITHOUT ghost (when gui false) — only real robot model
    rviz_node_no_ghost = Node(
        package='rviz2',
        executable='rviz2',
        name='rviz2',
        arguments=['-d', rviz_config_path],
        parameters=[
            {'robot_description': robot_description_content}
        ],
        output='screen',
        condition=UnlessCondition(LaunchConfiguration('gui'))
    )

    # MODIFIED: Conditionally launch the GUI node using IfCondition.
    # This node will only be launched if the 'gui' launch argument is set to 'true'.
    joint_control_gui_node = Node(
        package='zoros',
        executable='joint_control_gui.py',
        name='joint_control_gui',
        output='screen',
        condition=IfCondition(LaunchConfiguration('gui'))
    )

    # --- 4. Assemble and Return the Launch Description ---
    
    return LaunchDescription([
        # Launch arguments
        model_arg,
        simulation_arg,
        gui_arg,

        # Nodes
        robot_manager_node,
        robot_state_publisher_node,
        static_tf_ghost,
        ghost_robot_state_publisher_node,
        rviz_node_with_ghost,
        rviz_node_no_ghost,
        joint_control_gui_node  # This node is now conditional
    ])
